{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from hloc import extract_features, match_features, reconstruction, visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "In this notebook, we will run SfM reconstruction from scratch on a set of images. We choose the [South-Building dataset](https://openaccess.thecvf.com/content_cvpr_2013/html/Hane_Joint_3D_Scene_2013_CVPR_paper.html) - we will download it later. First, we define some paths."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = Path('datasets/sfm_South-Building/')\n",
    "images = dataset / 'South-Building/images/'\n",
    "\n",
    "outputs = Path('outputs/sfm/')\n",
    "sfm_pairs = outputs / 'pairs-exhaustive.txt'  # exhaustive matching\n",
    "sfm_dir = outputs / 'sfm_superpoint+superglue'\n",
    "\n",
    "feature_conf = extract_features.confs['superpoint_aachen']\n",
    "matcher_conf = match_features.confs['superglue']\n",
    "\n",
    "features = feature_conf['output']\n",
    "feature_file = f\"{features}.h5\"\n",
    "match_file = f\"{features}_{matcher_conf['output']}_{sfm_pairs.stem}.h5\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the dataset\n",
    "The dataset is simply a set of images. The intrinsic parameters will be extracted from the EXIF data, and refined with SfM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s \"$dataset\"\n",
    "wget http://cvg.ethz.ch/research/local-feature-evaluation/South-Building.zip -P $1\n",
    "unzip datasets/South-Building.zip -d $1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract local features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extract_features.main(feature_conf, images, outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exhaustive matching\n",
    "Since the dataset is small, we can match all $\\frac{n(n-1)}{2}$ images pairs. To do so, we pass the argument `exhaustive=True` and make sure that the pair file does not exist yet. If your dataset is larger, exhaustive matching might take a long time - consider selecting fewer pairs using image retrieval and `hloc/pairs_from_retrieval.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "match_features.main(matcher_conf, sfm_pairs, features, outputs, exhaustive=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SfM reconstruction\n",
    "Run COLMAP on the features and matches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reconstruction.main(\n",
    "    sfm_dir,\n",
    "    images,\n",
    "    sfm_pairs,\n",
    "    outputs / feature_file,\n",
    "    outputs / match_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization\n",
    "We visualize some of the registered images, and color their keypoint by visibility, track length, or triangulated depth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualization.visualize_sfm_2d(sfm_dir / 'models/0', images, color_by='visibility', n=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualization.visualize_sfm_2d(sfm_dir / 'models/0', images, color_by='track_length', n=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualization.visualize_sfm_2d(sfm_dir / 'models/0', images, color_by='depth', n=5)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
